from enthought.mayavi import mlab
import scipy
import affichAt3d
import numpy as np

def Lieliaison3d(mole,Elm):
    """lier les liasons chimiques dans 3D"""
    res=[] #to stock the tuple (i,j) if two atoms are linked
    for i in range(mole.NbTa):
        for j in range(i):
            if affichAt3d.Iflier(mole,i,j,Elm):
                res.append((i,j))
    
    for i,j in res:
        x=[mole.atom[i].pos[0],mole.atom[j].pos[0]]
        y=[mole.atom[i].pos[1],mole.atom[j].pos[1]]
        z=[mole.atom[i].pos[2],mole.atom[j].pos[2]]
        mlab.plot3d(x,y,z,opacity=0.7,tube_radius=0.07)
        
def YaAtom(mole,Elm):
    global fig1
    fig1=mlab.figure(2,size=(550, 550))
    mlab.clf()
    
    Xatom=[]
    Yatom=[]
    Zatom=[]
    atlb=[] #atoms' labels
    for i in range(mole.NbTa):
        x=mole.atom[i].pos[0]
        y=mole.atom[i].pos[1]
        z=mole.atom[i].pos[2]
        Xatom.append(x)
        Yatom.append(y)
        Zatom.append(z)        
        element='  '+str(i+1)+Elm[mole.atom[i].nb]
        atlb.append(element)

        
    #mlab.text3d(0,0,-1, 'element')
    
    points = mlab.points3d(Xatom,Yatom,Zatom,scale_factor=0.8,
                           color=(.8, .8, .8),scale_mode='none')
        
    extentA = (mole.ori[0],mole.ori[0]+mole.nbpas[0]*mole.pas[0],
               mole.ori[1],mole.ori[1]+mole.nbpas[1]*mole.pas[1],
               mole.ori[2],mole.ori[2]+mole.nbpas[2]*mole.pas[2])
    
    mlab.outline(points, color=(.7, .7, .7),extent=extentA)
    Lieliaison3d(mole,Elm)
    mlab.axes(points, color=(.7, .7, .7), extent=extentA, nb_labels=5,
            xlabel='X', ylabel='Y',zlabel='Z')
    
    for i in range(mole.NbTa):
        mlab.text3d(Xatom[i],Yatom[i],Zatom[i],atlb[i],
                    scale=0.6,color=(0.9,0.9,0.8))
    #mlab.text3d(Xatom,Yatom,Zatom,atlb)
    
def Crossproduct(X1,X2): #Majuscule signifies a vector in 3D
    """mathmatic dot product for two vectors"""
    x=X1[1]*X2[2]-X1[2]*X2[1]
    y=X1[2]*X2[0]-X1[0]*X2[2]
    z=X1[0]*X2[1]-X1[1]*X2[0]
    return [x,y,z]   
        
       

def YaAfficheplan(XX,mole,Elm):
    '''
    show the plan defined in mayavi in 3D
    '''
    global fig1
    print XX
    Vplan=Crossproduct((np.array(XX[2])-np.array(XX[0])),(np.array(XX[1])-np.array(XX[0])))
    Vori=XX[0]
    
    summe=0
    for i in range(3):
        summe=summe+Vori[i]*Vplan[i]
      
    def f(x, y):
        return ((summe-Vplan[0]*x-Vplan[1]*y)/Vplan[2])
    
    
    d1=mole.nbpas[0]*mole.pas[0]/30
    d2=mole.nbpas[1]*mole.pas[1]/30
    x, y = np.mgrid[mole.ori[0]:(mole.ori[0]+mole.nbpas[0]*mole.pas[0]):d1,\
                       mole.ori[1]:(mole.ori[1]+mole.nbpas[1]*mole.pas[1]):d2]
    
    YaAtom(mole,Elm)
    mlab.surf(x, y, f ,opacity=0.5)
   

    
    
    
    